# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import enum
import getpass
import os
import sys
import queue
import threading

from coin.util.slurm.driver import SlurmDriver
from coin.util.slurm.task import SlurmTask
from coin.util.slurm.base import Future

import tornado.ioloop
import tornado.platform.asyncio


class SlurmExecutor(object):
  """Execute calls asynchronously on Slurm.

  It provides the exact same interface to concurrent.futures.ThreadPoolExecutor
  and ProcessPoolExecutor.

  ----------
  Example #1
  ----------
    from coin.util.slurm.executor import SlurmExecutor

    with SlurmExecutor(max_workers=1) as executor:
      future = executor.submit(pow, 323, 1235)
      print(future.result())


  ----------
  Example #2
  ----------
    from coin.util.slurm.executor import SlurmExecutor

    futs = []
    with SlurmExecutor(max_workers=30, spread_jobs=True) as executor:
      for i in range(100):
        future = executor.submit(func, i)
        futs.append(future)

    for future in futs:
      print(future.result())
  """
  class StopState(enum.Enum):
    NONE = 0
    REQ_STOP = 1
    REQ_STOP_IMMEDIATELY = 2
    IN_PROGRESS_OR_STOPPED = 3

  @staticmethod
  def initialize():  # Alias of SlurmDriver.initialize
    SlurmDriver.initialize()

  def __init__(self,
               max_workers=5,
               spread_job=False,
               slurm_user=None,
               slurm_grid=None,
               force_use_slurm_venv=True,
               slurm_venv=None,
               slurm_shared_tempdir=None):
    SlurmDriver.initialize()

    self._ioloop = None
    self._driver = None
    self._driver_thread = None

    self._driver_startup_future = Future()
    self._driver_future = Future()

    self._stop_state = SlurmExecutor.StopState.NONE
    self._task_idx = 0
    self._pending_tasks = queue.Queue()
    self._running_tasks = set()

    # Slurm driver arguments
    assert max_workers > 0
    self._max_workers = max_workers
    self._spread_jobs = spread_job

    # Slurm task arguments
    user = slurm_user or getpass.getuser()
    grid = slurm_grid or 'iosg'
    self._slurm_task_kwargs = {'user': user, 'grid': grid, 'shared_tempdir': slurm_shared_tempdir}

    self._slurm_venv = slurm_venv
    if force_use_slurm_venv and slurm_venv is None:
      self._slurm_venv = os.environ.get(
          'COIN_SLURM_VENV', '/remote/{grid}/home/{user}/venv'.format(grid=grid, user=user))

  @property
  def _driver_state(self):
    if self._driver is None:
      return SlurmDriver.State.INIT
    return self._driver.state

  def submit(self, func, *args, **kwargs):
    self._raise_on_driver_error()
    assert self._driver_state == SlurmDriver.State.READY

    task_id = self._task_idx
    slurm_task = SlurmTask(func,
                           args=args,
                           kwargs=kwargs,
                           task_file_prefix=self._driver.slurm_job_id,
                           **self._slurm_task_kwargs)
    future = Future()
    self._pending_tasks.put((task_id, slurm_task, future))
    self._ioloop.add_callback(self._run_next_task)

    self._task_idx += 1
    return future

  # TODO(jaewon): Implement this
  # def map(func, *iterables, timeout=None, chunksize=1):
  #   pass

  def shutdown(self, wait=True):
    if self._driver_thread is None:
      return
    self._raise_on_driver_error()

    if wait:
      self._stop_state = SlurmExecutor.StopState.REQ_STOP
    else:
      self._cancel_all_pending_tasks()
      self._stop_state = SlurmExecutor.StopState.REQ_STOP_IMMEDIATELY

    self._ioloop.add_callback(self._run_next_task)
    self._wait_for_driver_stop()

  def _cancel_all_pending_tasks(self):
    while True:
      try:
        _task_id, slurm_task, future = self._pending_tasks.get_nowait()
      except queue.Empty:
        break
      del slurm_task
      future.cancel()

  def __enter__(self):
    assert self._driver_thread is None
    self._driver_thread = threading.Thread(target=self._thread_loop)
    self._driver_thread.start()
    self._driver_startup_future.result()
    return self

  def __exit__(self, exc_type, exc_value, exc_traceback):
    self.shutdown(wait=True)
    return False

  def _raise_on_driver_error(self):
    if self._driver_state == SlurmDriver.State.ERROR:
      self._cancel_all_pending_tasks()
      self._wait_for_driver_stop()

  def _wait_for_driver_stop(self):
    try:
      self._driver_future.result()  # Will raise exception.
    finally:
      self._driver_thread.join()
      self._driver_thread = None

  def _on_driver_error(self, *exc_info):
    assert self._driver_state == SlurmDriver.State.ERROR
    self._driver_future.set_exception_info(exc_info)
    self._ioloop.add_callback(self._ioloop.stop)

  def _run_next_task(self):
    if self._driver_state != SlurmDriver.State.READY:
      return

    while len(self._running_tasks) < self._max_workers:
      try:
        task_id, slurm_task, future = self._pending_tasks.get_nowait()
      except queue.Empty:
        break
      if not future.set_running_or_notify_cancel():  # Canceled
        continue
      self._running_tasks.add(task_id)
      self._ioloop.add_callback(self._run_task, task_id, slurm_task, future)

    if self._pending_tasks.empty():
      if (self._stop_state == SlurmExecutor.StopState.REQ_STOP_IMMEDIATELY
          or (self._stop_state == SlurmExecutor.StopState.REQ_STOP and not self._running_tasks)):
        self._stop_state = SlurmExecutor.StopState.IN_PROGRESS_OR_STOPPED
        self._ioloop.add_callback(self._stop_driver)
      return

  async def _run_task(self, task_id, slurm_task, future):
    try:
      with slurm_task:
        srun_args = [
            '-N1', '-n1', '--exclusive', *slurm_task.get_command_line(venv=self._slurm_venv)
        ]
        await self._driver.srun(srun_args)
        slurm_task.read_output_to_future(future)
        del slurm_task
    except Exception:
      future.set_exception_info(sys.exc_info())
    finally:
      self._running_tasks.remove(task_id)
      self._run_next_task()

  async def _start_driver(self):
    try:
      salloc_argv = ['-n', str(self._max_workers)]
      if self._spread_jobs:
        salloc_argv += ['--spread-job']

      self._driver = SlurmDriver()
      await self._driver.start(salloc_argv=salloc_argv, error_callback=self._on_driver_error)
      self._driver_startup_future.set_result(None)
    except Exception:
      self._driver_startup_future.set_exception_info(sys.exc_info())

  async def _stop_driver(self):
    try:
      await self._driver.stop()
      self._driver_future.set_result(None)
    except Exception:
      self._driver_future.set_exception_info(sys.exc_info())
    finally:
      self._ioloop.add_callback(self._ioloop.stop)

  def _thread_loop(self):
    self._ioloop = tornado.ioloop.IOLoop.current()
    self._ioloop.add_callback(self._start_driver)
    self._ioloop.start()
